import os
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE


def create_save_dir(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)


def add_noise(images, noise_level=0.1):
    """为图像添加噪声"""
    noisy_images = images + noise_level * np.random.normal(size=images.shape)
    return np.clip(noisy_images, 0, 1)  # 保证数据在[0, 1]范围


def apply_pca(images):
    """应用PCA降维并返回降维后的数据和重构数据"""
    pca = PCA(n_components=0.95)
    images_pca = pca.fit_transform(images)
    images_reconstructed = pca.inverse_transform(images_pca)
    return images_pca, images_reconstructed


def save_class_images(images, noisy_images, denoised_images, labels, class_names):
    """保存每个类别的原始、噪声和降噪图像"""
    save_dir = 'images/class_0.95/'
    create_save_dir(save_dir)

    for class_index in range(10):
        # 获取每个类别的索引
        class_indices = np.where(labels.numpy() == class_index)[0]

        # 如果这个类别有图像
        if len(class_indices) > 0:
            # 只输出第一个图像
            index = class_indices[0]

            plt.figure(figsize=(15, 5))

            # 保存原始图像
            plt.subplot(1, 3, 1)
            plt.imshow(images[index].transpose((1, 2, 0)))
            plt.title(f"Original - {class_names[class_index]}")
            plt.axis('off')

            # 保存噪声图像
            plt.subplot(1, 3, 2)
            noisy_image = noisy_images[index].reshape(3, 32, 32)  # 将一维数据重新reshape回[3, 32, 32]
            plt.imshow(noisy_image.transpose((1, 2, 0)))  # 转换维度以匹配imshow的要求
            plt.title(f"Noisy - {class_names[class_index]}")
            plt.axis('off')

            # 保存降噪后的图像
            plt.subplot(1, 3, 3)
            plt.imshow(denoised_images[index].transpose((1, 2, 0)))
            plt.title(f"Denoised - {class_names[class_index]}")
            plt.axis('off')

            plt.suptitle(f'Class {class_index} - {class_names[class_index]}')

            # 保存图片到文件
            plt.savefig(f'{save_dir}Class_{class_index}_0.png')
            plt.show()


def plot_visualization(images, labels, title, method, perplexity=30):
    """绘制PCA或t-SNE的可视化"""
    if method == 'PCA':
        model = PCA(n_components=2)
    elif method == 't-SNE':
        # 如果使用t-SNE，确保perplexity小于样本数量
        model = TSNE(n_components=2, random_state=42, perplexity=min(perplexity, max(5, images.size(0) - 1)))
    else:
        raise ValueError("Method must be either 'PCA' or 't-SNE'")

    images_flattened = images.view(images.size(0), -1).numpy()
    X_reduced = model.fit_transform(images_flattened)

    plt.figure(figsize=(10, 8))
    scatter = plt.scatter(X_reduced[:, 0], X_reduced[:, 1], c=labels.numpy(), cmap='viridis', alpha=0.7, edgecolor='k')
    plt.colorbar(scatter)
    plt.xlabel('Component 1' if method == 'PCA' else 't-SNE Feature 1')
    plt.ylabel('Component 2' if method == 'PCA' else 't-SNE Feature 2')
    plt.title(title)

    # 保存图像到文件夹
    create_save_dir('images')
    plt.savefig(f'images/{title}-{method}.png')
    plt.show()


def main():
    # 定义数据转换
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # 加载测试数据
    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

    # 增大 batch_size 以确保有足够的样本进行降维和可视化
    batch_size = 10000
    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)

    # 获取一批测试数据
    dataiter = iter(testloader)
    images, labels = next(dataiter)

    # 检查实际获取的样本数量
    actual_batch_size = images.size(0)
    # print(f"实际批量大小: {actual_batch_size}")

    # 将图像数据转换为一维
    X_cifar10_test = images.view(images.size(0), -1).numpy()

    # 为CIFAR-10数据添加噪声
    noisy_cifar10_test = add_noise(X_cifar10_test)

    # 使用PCA对噪声数据进行降维和重构
    _, X_pca_inverse_test = apply_pca(noisy_cifar10_test)
    X_pca_inverse_test = X_pca_inverse_test.reshape(-1, 3, 32, 32)
    X_pca_inverse_test = torch.tensor(X_pca_inverse_test)

    # 保存每个类别的图像（原始、噪声、降噪）
    save_class_images(images.numpy(), noisy_cifar10_test, X_pca_inverse_test.numpy(), labels, testset.classes)

    # 使用PCA和t-SNE进行可视化
    plot_visualization(images, labels, 'CIFAR-10 Original', 'PCA')
    plot_visualization(images, labels, 'CIFAR-10 Original', 't-SNE', perplexity=50)  # 确保 perplexity < batch_size
    plot_visualization(torch.tensor(noisy_cifar10_test), labels, 'CIFAR-10 Noisy', 'PCA')
    plot_visualization(torch.tensor(noisy_cifar10_test), labels, 'CIFAR-10 Noisy', 't-SNE', perplexity=50)
    plot_visualization(X_pca_inverse_test, labels, 'CIFAR-10 Denoised', 'PCA')
    plot_visualization(X_pca_inverse_test, labels, 'CIFAR-10 Denoised', 't-SNE', perplexity=50)


if __name__ == '__main__':
    main()
